# Copyright (c) 2017 Tencent Inc.
# All rights reserved.
#
# Author: Li Wenting <wentingli@tencent.com>
# Date:   October 27, 2017

"""
This module defines various build functions for building
targets from sources and custom parameters.
"""


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import fnmatch
import getpass
import os
import shutil
import socket
import sys
import tarfile
import textwrap
import traceback
import time
import zipfile

from blade import console
from blade import util


# These following helper functions is designed to centralize error handling


_outputs = []


def _declare_outputs(*outputs):
    """Declare output files, which can be verified after ran and cleanup on error."""
    _outputs[:] = outputs


def _verify_outputs():
    """Verify whether the declared output files were correctly generated."""
    missing = [o for o in _outputs if not os.path.exists(o)]
    if missing:
        raise FileNotFoundError('"%s" is not generated' % missing)


def _cleanup_outputs():
    """Cleanup declared output files on error."""
    console.debug('Cleanup:%s' % _outputs)
    for output in _outputs:
        try:
            os.remove(output)
        except OSError:
            pass


# The build functions are defined as follows form:
#
#    def build_function_name(kwargs..., args):
#        pass
#
#    Args:
#        * kwargs...: name=value pairs as parameters in command line
#        * args: any other non-kw args
#
#    Return:
#        Return None(include implicit return) or 0 on success, otherwise return
#        a positive int (e.g. 1) or raise an exception to indicate failure.
#
#    When call this from the command line, all arguments which match the
#    `--name=value` pattern will be converted into kwargs, any other arguments
#    are merged into the `args` argument. So any `name` must be a valid python
#    identifier.


def generate_scm(scm, revision, url, profile, compiler, args):
    """Generate `scm.c` file"""

    _declare_outputs(scm)

    version = '%s@%s' % (url, revision)
    with open(scm, 'w') as f:
        f.write(textwrap.dedent(r'''\
                /* This file was generated by blade */
                extern "C" {
                namespace binary_version {
                extern const int kSvnInfoCount = 1;
                extern const char* const kSvnInfo[] = {"%s\n"};
                extern const int kScmInfoCount = 1;
                extern const char* const kScmInfo[] = {"%s\n"};
                extern const char kBuildType[] = "%s";
                extern const char kBuildTime[] = "%s";
                extern const char kBuilderName[] = "%s";
                extern const char kHostName[] = "%s";
                extern const char kCompiler[] = "%s";
                }}''') % (version,
                          version,
                          profile,
                          time.asctime(),
                          getpass.getuser(),
                          socket.gethostname(),
                          compiler))


def generate_cc_inclusion_check(args):
    # Import from function to avoid affecting the performance of other tools.
    from blade import inclusion_check  # pylint: disable=import-outside-toplevel
    result_file = args[0]
    info_file = args[1]
    _declare_outputs(result_file)
    console.set_log_file('%s.log' % info_file)
    console.enable_color(True)
    ok, details = inclusion_check.check(info_file)
    with open(info_file + '.details', 'wb') as f:
        util.pickle.dump(details, f)
    if not ok:
        return 1
    with open(result_file, 'w') as f:
        f.write('OK')
    return None


_PACKAGE_MANIFEST = 'MANIFEST.TXT'


def archive_package_sources(package, sources, destinations):
    """Content of the `MANIFEST.TXT` file in the target zip file"""
    manifest = []
    for i, s in enumerate(sources):
        package(s, destinations[i])
        manifest.append('%s %s' % (util.md5sum_file(s), destinations[i]))
    return manifest


def generate_zip_package(path, sources, destinations):
    zip = zipfile.ZipFile(path, 'w', zipfile.ZIP_DEFLATED, allowZip64=True)
    manifest = archive_package_sources(zip.write, sources, destinations)
    zip.writestr(_PACKAGE_MANIFEST, '\n'.join(manifest) + '\n')
    zip.close()


def generate_tar_package(path, sources, destinations, mode):
    tar = tarfile.open(path, mode, dereference=True)
    manifest = archive_package_sources(tar.add, sources, destinations)
    manifest_path = '%s.MANIFEST' % path
    m = open(manifest_path, 'w')
    m.write('\n'.join(manifest) + '\n\n')
    m.close()
    tar.add(manifest_path, _PACKAGE_MANIFEST)
    tar.close()


_TAR_WRITE_MODES = {
    'tar': 'w',
    'tar.gz': 'w:gz',
    'tgz': 'w:gz',
    'tar.bz2': 'w:bz2',
    'tbz': 'w:bz2',
}


def _tar_write_mode(path):
    """Get the tar write mode from extname of path"""
    for ext, mode in _TAR_WRITE_MODES.items():
        if path.endswith(ext):
            return mode
    return ''


def generate_package(args):
    path = args[0]
    manifest = args[1:]
    _declare_outputs(path)
    assert len(manifest) % 2 == 0
    middle = len(manifest) // 2
    sources = manifest[:middle]
    destinations = manifest[middle:]
    if path.endswith('.zip'):
        generate_zip_package(path, sources, destinations)
    else:
        mode = _tar_write_mode(path)
        generate_tar_package(path, sources, destinations, mode)


def _generate_resource_index(targets, sources, name, path):
    """Generate resource index description file for a cc resource library"""
    header, source = targets
    with open(header, 'w') as h, open(source, 'w') as c:
        full_name = util.regular_variable_name(os.path.join(path, name))
        guard_name = 'BLADE_RESOURCE_%s_H_' % full_name.upper()
        index_name = 'RESOURCE_INDEX_%s' % full_name

        h.write(textwrap.dedent('''\
                // This file was automatically generated by blade
                #ifndef {0}
                #define {0}

                #ifdef __cplusplus
                extern "C" {{
                #endif

                #ifndef BLADE_RESOURCE_TYPE_DEFINED
                #define BLADE_RESOURCE_TYPE_DEFINED
                struct BladeResourceEntry {{
                    const char* name;
                    const char* data;
                    unsigned int size;
                }};
                #endif''').format(guard_name))
        c.write(textwrap.dedent('''\
                // This file was automatically generated by blade
                #include "{0}"

                const struct BladeResourceEntry {1}[] = {{''').format(header, index_name))

        for s in sources:
            entry_var = util.regular_variable_name(s)
            entry_name = os.path.relpath(s, path)
            entry_size = os.path.getsize(s)
            h.write('// %s\n' % entry_name)
            h.write('extern const char RESOURCE_%s[%d];\n' % (entry_var, entry_size))
            h.write('extern const unsigned RESOURCE_%s_len;\n' % entry_var)
            c.write('    { "%s", RESOURCE_%s, %s },\n' % (entry_name, entry_var, entry_size))

        c.write(textwrap.dedent('''\
                }};
                const unsigned {0}_len = {1};''').format(index_name, len(sources)))
        h.write(textwrap.dedent('''\
                // Resource index
                extern const struct BladeResourceEntry {0}[];
                extern const unsigned {0}_len;

                #ifdef __cplusplus
                }}  // extern "C"
                #endif

                #endif  // {1}''').format(index_name, guard_name))


def generate_resource_index(args):
    name, path = args[0], args[1]
    targets = args[2], args[3]
    sources = args[4:]
    _declare_outputs(*targets)
    return _generate_resource_index(targets, sources, name, path)


def generate_java_jar(compression_level, args):
    jar, target = args[0], args[1]
    _declare_outputs(target)
    resources_dir = target.replace('.jar', '.resources')
    arg = args[2]
    if arg.endswith('__classes__.jar'):
        classes_jar = arg
        resources = args[3:]
    else:
        classes_jar = ''
        resources = args[2:]

    def archive_resources(resources_dir, resources, new=True):
        options = ('c' if new else 'u') + 'f' + compression_level
        cmd = ['%s %s %s' % (jar, options, target)]
        for resource in resources:
            cmd.append("-C '%s' '%s'" % (resources_dir,
                                         os.path.relpath(resource, resources_dir)))
        return util.shell(cmd)

    if classes_jar:
        shutil.copy2(classes_jar, target)
        if resources:
            return archive_resources(resources_dir, resources, False)
    else:
        return archive_resources(resources_dir, resources, True)


def generate_java_resource(args):
    assert len(args) % 2 == 0
    middle = len(args) // 2
    targets = args[:middle]
    sources = args[middle:]
    _declare_outputs(*targets)
    for i in range(middle):
        shutil.copy(sources[i], targets[i])


def _get_all_test_class_names_in_jar(jar):
    """Returns a list of test class names in the jar file."""
    test_class_names = []
    zip_file = zipfile.ZipFile(jar, 'r')
    name_list = zip_file.namelist()
    for name in name_list:
        basename = os.path.basename(name)
        # Exclude inner class and Test.class
        if (basename.endswith('Test.class') and
                len(basename) > len('Test.class') and '$' not in basename):
            class_name = name.replace('/', '.')[:-6]  # Remove .class suffix
            test_class_names.append(class_name)
    zip_file.close()
    return test_class_names


def _jacoco_test_coverage_flag(jacocoagent, packages_under_test):
    if packages_under_test and jacocoagent:
        jacocoagent = os.path.abspath(jacocoagent)
        packages = packages_under_test.split(':')
        options = [
            'includes=%s' % ':'.join([p + '.*' for p in packages if p]),
            'output=file',
        ]
        return '-javaagent:%s=%s' % (jacocoagent, ','.join(options))
    return ''


def generate_java_test(script, main_class, jacocoagent, packages_under_test, args):
    _declare_outputs(script)
    jars = args
    test_jar = jars[0]
    test_classes = ' '.join(_get_all_test_class_names_in_jar(test_jar))
    with open(script, 'w') as f:
        coverage_flags = _jacoco_test_coverage_flag(jacocoagent, packages_under_test)
        f.write(textwrap.dedent('''\
                #!/bin/sh
                # Auto generated wrapper shell script by blade

                if [ -n "$BLADE_COVERAGE" ]; then
                    coverage_options="%s"
                fi

                exec java $coverage_options -classpath %s %s %s $@''') % (
                coverage_flags, ':'.join(jars), main_class, test_classes))
    os.chmod(script, 0o755)


def generate_fat_jar(output, **kwargs):
    # Import from function to avoid affecting the performance of other tools.
    from blade import fatjar  # pylint: disable=import-outside-toplevel
    _declare_outputs(output)
    console.set_log_file('%s.log' % output.replace('.fat.jar', '__fatjar__'))
    console.enable_color(True)
    fatjar.generate_fat_jar(output=output, **kwargs)


def generate_one_jar(onejar, main_class, bootjar, args):
    _declare_outputs(onejar)
    # Assume the first jar is the main jar, others jars are dependencies.
    main_jar = args[0]
    jars = args[1:]
    path = onejar
    onejar = zipfile.ZipFile(path, 'w', allowZip64=True)
    jar_path_set = set()
    # Copy files from one-jar-boot.jar to the target jar
    zip_file = zipfile.ZipFile(bootjar, 'r')
    name_list = zip_file.namelist()
    for name in name_list:
        if not name.lower().endswith('manifest.mf'):  # Exclude manifest
            onejar.writestr(name, zip_file.read(name))
            jar_path_set.add(name)
    zip_file.close()

    # Main jar and dependencies
    onejar.write(main_jar, os.path.join('main',
                                        os.path.basename(main_jar)))
    for dep in jars:
        dep_name = os.path.basename(dep)
        onejar.write(dep, os.path.join('lib', dep_name))

    # Copy resources to the root of target onejar
    for jar in [main_jar] + jars:
        jar = zipfile.ZipFile(jar, 'r')
        jar_name_list = jar.namelist()
        for name in jar_name_list:
            if name.endswith('.class') or name.upper().startswith('META-INF'):
                continue
            if name not in jar_path_set:
                jar_path_set.add(name)
                onejar.writestr(name, jar.read(name))
        jar.close()

    # Manifest
    # Note that the manifest file must end with a new line or carriage return
    onejar.writestr(os.path.join('META-INF', 'MANIFEST.MF'),
                    textwrap.dedent('''\
                            Manifest-Version: 1.0
                            Main-Class: com.simontuffs.onejar.Boot
                            One-Jar-Main-Class: %s

                            ''') % main_class)
    onejar.close()


def generate_java_binary(args):
    script, onejar = args
    _declare_outputs(script)
    basename = os.path.basename(onejar)
    fullpath = os.path.abspath(onejar)
    with open(script, 'w') as f:
        f.write(textwrap.dedent('''\
                #!/bin/sh
                # Auto generated wrapper shell script by blade

                jar=`dirname "$0"`/"%s"
                if [ ! -f "$jar" ]; then
                  jar="%s"
                fi

                exec java -jar "$jar" $@
                ''') % (basename, fullpath))
    os.chmod(script, 0o755)


def generate_scala_test(script, java, scala, jacocoagent, packages_under_test, args):
    _declare_outputs(script)
    jars = args
    test_jar = jars[0]
    test_class_names = _get_all_test_class_names_in_jar(test_jar)
    scala, java = os.path.abspath(scala), os.path.abspath(java)
    java_args = ''
    coverage_flags = _jacoco_test_coverage_flag(jacocoagent, packages_under_test)
    if coverage_flags:
        java_args = '-J%s' % coverage_flags
    run_args = 'org.scalatest.run ' + ' '.join(test_class_names)
    with open(script, 'w') as f:
        text = textwrap.dedent('''\
                #!/bin/sh
                # Auto generated wrapper shell script by blade

                if [ -n "$BLADE_COVERAGE" ]; then
                    coverage_options="%s"
                fi

                JAVACMD=%s exec %s "$coverage_options" -classpath %s %s $@
                ''') % (java_args, java, scala, ':'.join(jars), run_args)
        f.write(text)
    os.chmod(script, 0o755)


def generate_shell_test(args):
    wrapper = args[0]
    scripts = args[1:]
    _declare_outputs(wrapper)
    with open(wrapper, 'w') as f:
        f.write(textwrap.dedent("""\
                #!/bin/sh
                # Auto generated wrapper shell script by blade

                set -e

                %s

                """) % '\n'.join(['. %s' % os.path.abspath(s) for s in scripts]))
    os.chmod(wrapper, 0o755)


def generate_shell_testdata(args):
    path = args[0]
    testdata = args[1:]
    _declare_outputs(path)
    assert len(testdata) % 2 == 0
    middle = len(testdata) // 2
    sources = testdata[:middle]
    destinations = testdata[middle:]
    with open(path, 'w') as f:
        for i in range(middle):
            f.write('%s %s\n' % (os.path.abspath(sources[i]), destinations[i]))


def generate_python_library(pylib, basedir, args):
    _declare_outputs(pylib)
    sources = []
    for py in args:
        digest = util.md5sum_file(py)
        sources.append((py, digest))
    with open(pylib, 'w') as f:
        print(str({
            'base_dir': basedir,
            'srcs': sources
        }), file=f)


def _is_python_excluded_path(filename, exclusions):
    for exclusion in exclusions:
        if fnmatch.fnmatch(filename, exclusion):
            return True
    return False


def _update_init_py_dirs(arcname, dirs, dirs_with_init_py):
    dir = os.path.dirname(arcname)
    if os.path.basename(arcname) == '__init__.py':
        dirs_with_init_py.add(dir)
    while dir:
        dirs.add(dir)
        dir = os.path.dirname(dir)


def _pybin_add_pylib(pybin, libname, exclusions, dirs, dirs_with_init_py):
    with open(libname) as pylib:
        data = eval(pylib.read())  # pylint: disable=eval-used
        pylib_base_dir = data['base_dir']
        for libsrc, digest in data['srcs']:
            arcname = os.path.relpath(libsrc, pylib_base_dir)
            if not _is_python_excluded_path(arcname, exclusions):
                _update_init_py_dirs(arcname, dirs, dirs_with_init_py)
                pybin.write(libsrc, arcname)


def _pybin_add_zip(pybin, libname, filter, exclusions, dirs, dirs_with_init_py):
    with zipfile.ZipFile(libname, 'r') as lib:
        name_list = lib.namelist()
        for name in name_list:
            if filter(name) and not _is_python_excluded_path(name, exclusions):
                if dirs is not None and dirs_with_init_py is not None:
                    _update_init_py_dirs(name, dirs, dirs_with_init_py)
                pybin.writestr(name, lib.read(name))


def _pybin_add_egg(pybin, libname, exclusions):
    def filter(name):
        if name.startswith('EGG-INFO'):
            return False
        if name.endswith('.pyc'):
            return False
        return True

    _pybin_add_zip(pybin, libname, filter, exclusions, None, None)


def _pybin_add_whl(pybin, libname, exclusions, dirs, dirs_with_init_py):
    def filter(name):
        if '.dist-info/' in name:
            return False
        return True

    _pybin_add_zip(pybin, libname, filter, exclusions, dirs, dirs_with_init_py)


def generate_python_binary(pybin, basedir, exclusions, mainentry, args):
    _declare_outputs(pybin)
    pybin_zip = zipfile.ZipFile(pybin, 'w', zipfile.ZIP_DEFLATED, allowZip64=True)
    exclusions = exclusions.split(',')
    dirs, dirs_with_init_py = set(), set()
    for arg in args:
        if arg.endswith('.pylib'):
            _pybin_add_pylib(pybin_zip, arg, exclusions, dirs, dirs_with_init_py)
        elif arg.endswith('.egg'):
            _pybin_add_egg(pybin_zip, arg, exclusions)
        elif arg.endswith('.whl'):
            _pybin_add_whl(pybin_zip, arg, exclusions, dirs, dirs_with_init_py)
        else:
            assert False, 'Unknown file type "%s" to build python_binary' % arg

    # Insert __init__.py into each dir if missing
    dirs_missing_init_py = dirs - dirs_with_init_py
    for dir in sorted(dirs_missing_init_py):
        pybin_zip.writestr(os.path.join(dir, '__init__.py'), '')
    pybin_zip.writestr('__init__.py', '')
    pybin_zip.close()

    with open(pybin, 'rb') as f:
        zip_content = f.read()
    # Insert bootstrap before zip, it is also a valid zip file.
    # unzip will seek actually start until meet the zip magic number.
    bootstrap = ('#!/bin/sh\n\n'
                 'PYTHONPATH="$0:$PYTHONPATH" exec python -m "%s" "$@"\n') % mainentry
    with open(pybin, 'wb') as f:
        f.write(bootstrap.encode('utf-8'))
        f.write(zip_content)
    os.chmod(pybin, 0o755)


_BUILTIN_TOOLS = {
    'scm': generate_scm,
    'package': generate_package,
    'cc_inclusion_check': generate_cc_inclusion_check,
    'resource_index': generate_resource_index,
    'java_jar': generate_java_jar,
    'java_resource': generate_java_resource,
    'java_test': generate_java_test,
    'java_fatjar': generate_fat_jar,
    'java_onejar': generate_one_jar,
    'java_binary': generate_java_binary,
    'scala_test': generate_scala_test,
    'shell_test': generate_shell_test,
    'shell_testdata': generate_shell_testdata,
    'python_library': generate_python_library,
    'python_binary': generate_python_binary,
}


def main():
    name = sys.argv[1]
    exit_code = 0
    try:
        options, args = util.parse_command_line(sys.argv[2:])
        exit_code = _BUILTIN_TOOLS[name](args=args, **options)
        if not exit_code:
            _verify_outputs()
    except Exception as e:  # pylint: disable=broad-except
        console.error('Blade build tool %s error: %s %s' % (name, str(e), traceback.format_exc()))
        exit_code = 1
    finally:
        if exit_code:
            _cleanup_outputs()
        sys.exit(exit_code or 0)


if __name__ == '__main__':
    main()
